# @日期：2021/1/17 12:15
# @公司：阿尔法工场
# @作者：张大鹏
# @昵称：lxgzhw
# @文件：02.保存加载和使用模型.py
# @描述：保存加载和使用模型
import torch
from torch import nn
import numpy as np

# 创建线性模型
# 参数1：输入的特征长度（一个受教育年限数据）
# 参数2：输出的特征长度（一个收入预测数据）
model = nn.Linear(1, 1)

# 保存模型
# torch.save(model.state_dict(), "EducationAndIncome.model")
# 加载模型
model.load_state_dict(torch.load("EducationAndIncome.model"))
model.eval()

# 使用模型预测数据
# 要求1：预测的数据必须是二维数组
# 要求2：预测的数据类型必须是np.float32
x = np.array([[10]]).astype(np.float32)
# 要求3：预测的数据必须转换为torch张量类型
x = torch.from_numpy(x)
prediction = model(x)
# print(prediction)
# 提取预测结果
# 注意：这里不能直接用prediction.numpy()
prediction = prediction.detach().numpy()
prediction = prediction[0][0]
print(prediction)
